
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include <fstream>
#include <string>
#include <vector>
using namespace std;

#include "Socket.h"
#include "dump.h"

#if defined(_WIN32) || defined(WIN32)
#include <process.h>
#endif

#pragma pack(1)

typedef struct
{
    unsigned short int      id;
    unsigned short int      flag;
    unsigned short int      qdcount;
    unsigned short int      ancount;
    unsigned short int      nscount;
    unsigned short int      arcount;
}DNSHEADER;

#if 0
int domain2labels(const char * domain, void * labels, int length)
{/*{{{*/
    char *          p = (char *)domain;
    unsigned char * result = (unsigned char *)labels;
    unsigned int    len;
    if(length < (int)strlen(domain) + 2)
        return 0;

    while(1)
    {
        char *  p1 = strchr(p, '.');
        if(p1 != NULL)
            len = p1 - p;
        else
            len = strlen(p);
        *result  = (unsigned char)(len & 0xff);
        memcpy(result + 1, p, len);
        result += len + 1;
        if(p1 == NULL)
            break;
        p = p1 + 1;
    }
    *result = '\0';
    return (result - (unsigned char *)labels + 1);
}/*}}}*/

void parsePacket(const void * packet, int len, const char * dns)
{/*{{{*/
    DNSHEADER *     header = (DNSHEADER *)packet;
    int     rcode = (header->flag & 0xf00);
    rcode >>= 8;
    char    ipaddr[32];
    char *  ip;

    switch(rcode)
    {
    case 0:
        // 正常识别。
        {
            const char * answers = (const char *)packet + sizeof(header);
            while(*answers != '\0')
                answers ++;
            answers += 5;
            ip = (char *)packet + len - 4;
            sprintf(ipaddr, "%d.%d.%d.%d", (ip[0] & 0xff), (ip[1] & 0xff), 
                ip[2] & 0xff, ip[3] & 0xff);
            PRINTF("[ OK ] result from %s: ip = %s\n", dns, ipaddr);
        }
        break;
    case 3:
        PRINTF("[FAIL] result from %s: no such name.\n", dns);
        // no such name.
        break;
    case 5:
        // refused.
        PRINTF("[FAIL] result from %s: refused.\n", dns);
        break;
    default:
        PRINTF("[FAIL] result from %s: rcode = %d\n", dns, rcode);
        break;
    }
}/*}}}*/


void * _detectproc(void * arg)
{/*{{{*/
    PRINTF("thread %d started!\n", (int)arg);
    ThreadRef();

    CSocket         sock(CSocket::eSocketType_UDP);
    DNSHEADER       header;
    char            packet[1024];
    unsigned char   question[1024];
    int             ret;
    int             dmlen;
    int             len;
    fd_set          rfds;
    struct timeval  tv;

    dmlen = domain2labels(g_domain.c_str(), question, sizeof(question));
    if(dmlen <= 0)
    {
        PRINTF("[FAIL] convert %s error.\n", g_domain.c_str());
        ThreadUnRef();
        return NULL;
    }

    sock.Create();
    while(1)
    {
        string  dns = _getnexthost();
        if(dns.length() == 0)
            break;
        memset(&header, 0, sizeof(header));
        header.id       = rand() % 0xffff;
        header.flag     = 1;
        header.qdcount  = htons(1);

        memcpy(packet, &header, sizeof(header));
        memcpy(packet + sizeof(header), question, dmlen);
        len = dmlen + sizeof(header);
        packet[len] = 0;
        packet[len + 1] = 1;
        packet[len + 2] = 0;
        packet[len + 3] = 1;
        len += 4;

        sock.SendTo(dns.c_str(), 53, packet, len);

        FD_ZERO(&rfds);
        FD_SET(sock.GetSocket(), &rfds);
        tv.tv_sec = 10;
        tv.tv_usec = 0;
        ret = select(sock.GetSocket() + 1, &rfds, NULL, NULL, &tv);
        if(ret < 0)
        {
            perror("select");
            continue;
        } else if(ret == 0) {
            PRINTF("[FAIL] query %s timeout.\n", dns.c_str());
        } else {
            ret = sock.Recv(packet, 1024);
            if(ret <= 0)
            {
                perror("recv");
                continue;
            }
            parsePacket(packet, ret, dns.c_str());
        }
    }
    sock.Close();
    ThreadUnRef();
    return NULL;
}/*}}}*/
#endif

int main(int argc, char * argv[])
{
    CSocket     sock(CSocket::eSocketType_UDP);
    char        packet[4096];

    try {
        sock.Create();
        sock.Bind(53);
        while(true) {
            int ret = sock.Recv(packet, 4096);
            dump(packet, ret);
            DNSHEADER *     header = (DNSHEADER *)packet;
            header->flag = 0x8281;
            sock.Send(packet, ret);
        }
    } catch(const char * e) {
        printf("%s\n", e);
    }
    return 0;
}



